import os
import sys
import json
import zlib
import pickle

sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))


import yaml
import paddle
import paddle.distributed as dist

paddle.seed(2)

from ppocr.data import build_dataloader
from ppocr.modeling.architectures import build_model
from ppocr.losses import build_loss
from ppocr.optimizer import build_optimizer
from ppocr.postprocess import build_post_process
from ppocr.metrics import build_metric
from ppocr.utils.save_load import init_model
import tools.program as program


def load_compressed_file(fname):
    with open(fname, 'rb') as fh:
        data = pickle.loads(zlib.decompress(fh.read()))
        
    return data

def save_compress_file(fname, data):
    with open(fname, 'wb') as fh:
        fh.write(zlib.compress(pickle.dumps(data)))


def main(config, device, logger):
    global_config = config['Global']
    # valid_dataloader = build_dataloader(config, 'Eval', device, logger)
    post_process_class = build_post_process(config['PostProcess'], global_config)

    if hasattr(post_process_class, 'character'):
        char_num = len(getattr(post_process_class, 'character'))
        config['Architecture']["Head"]['out_channels'] = char_num

    model = build_model(config['Architecture'])
    eval_class = build_metric(config['Metric'])
    pre_best_model_dict = init_model(config, model, logger, None)

    model.eval()
    # x = paddle.rand([1, 3, 32, 320])
    x = paddle.rand([1, 3, 320, 320])
    with paddle.no_grad():
        y = model(x)
    print(y['maps'].shape)



    # bboxes_results, pred_results = program.detect_bboxes(model, valid_dataloader, post_process_class, eval_class)

    # save_compress_file('output/pred_bboxes', [bboxes_results, pred_results])
        

    # metric = eval_class.get_metric()
    # print(metric)


if __name__ == '__main__':
    config, device, logger, vdl_writer = program.preprocess(is_train=False)
    main(config, device, logger)